
import numpy as np


from data_config import DataConfig
from dataset.CD_dataset import CDDataset
from mindspore import dataset as ds
import multiprocessing

def get_loader(data_name, img_size=256, batch_size=8, split='test',
               is_train=False, dataset='CDDataset'):
    dataConfig = DataConfig().get_data_config(data_name)
    root_dir = dataConfig.root_dir
    label_transform = dataConfig.label_transform

    if dataset == 'CDDataset':
        data_set = CDDataset(root_dir=root_dir, split=split,
                                 img_size=img_size, is_train=is_train,
                                 label_transform=label_transform)
    else:
        raise NotImplementedError(
            'Wrong dataset name %s (choose one from [CDDataset])'
            % dataset)
    data_loader = ds.GeneratorDataset(data_set, ['imgAB', 'label'], shuffle=is_train)
    data_loader = data_loader.batch(batch_size)

    return data_loader


def get_loaders(data_name='LEVIR', dataset='CDDataset', split='train', batch_size=4, img_size=256):
    data_name = data_name
    dataConfig = DataConfig().get_data_config(data_name)
    root_dir = dataConfig.root_dir
    label_transform = dataConfig.label_transform
    split = split
    split_val = 'val'

    if dataset == 'CDDataset':
        train_set = CDDataset(root_dir=root_dir, split=split,
                                 img_size=img_size, is_train=True,
                                 label_transform=label_transform)
        val_set = CDDataset(root_dir=root_dir, split=split_val,
                            img_size=img_size, is_train=False,
                            label_transform=label_transform)
    else:
        raise NotImplementedError(
            'Wrong dataset name %s (choose one from [CDDataset,])'
            % dataset)



    train_data_loader = ds.GeneratorDataset(train_set, ['imgAB', 'label'], shuffle=True)
    train_data_loader = train_data_loader.batch(batch_size)

    val_data_loader = ds.GeneratorDataset(val_set, ['imgAB', 'label'], shuffle=False)
    val_data_loader = val_data_loader.batch(1)

    return {'train': train_data_loader, 'val': val_data_loader}


def get_loader_multi_gpu(data_name='LEVIR', dataset='CDDataset', split='train', batch_size=4, img_size=256, num_shards=1, shard_id=0):
    data_name = data_name
    dataConfig = DataConfig().get_data_config(data_name)
    root_dir = dataConfig.root_dir
    label_transform = dataConfig.label_transform
    split = split
    
    ds.config.set_enable_shared_mem(True)
    cores = multiprocessing.cpu_count()
    num_parallel_workers = min(4, cores // num_shards)

    print("dataset use current rank id: {}/{} ".format(shard_id, num_shards))
    if dataset == 'CDDataset':
        train_set = CDDataset(root_dir=root_dir, split=split,
                                 img_size=img_size, is_train=True,
                                 label_transform=label_transform)
    else:
        raise NotImplementedError(
            'Wrong dataset name %s (choose one from [CDDataset,])'
            % dataset)

    # train_data_loader = ds.GeneratorDataset(train_set, ['imgAB', 'label'], shuffle=True)
    train_data_loader = ds.GeneratorDataset(train_set, ['imgAB', 'label'], shuffle=True,
                                            num_parallel_workers=1,
                                            num_shards=num_shards, shard_id=shard_id, python_multiprocessing=True)
    train_data_loader = train_data_loader.batch(batch_size)

    return train_data_loader

